import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph, END
from typing import TypedDict, List
import os
import multiprocessing
import time


class State(TypedDict):
    messages: List[str]
    step: int


def create_graph(checkpointer=None):
    """创建图实例"""

    def step1(state: State):
        print("Executing step1...")
        time.sleep(1)  # 模拟工作
        return {"messages": [*state["messages"], "Step 1 completed"], "step": 1}

    def step2(state: State):
        print("Executing step2...")
        time.sleep(1)  # 模拟工作
        return {"messages": [*state["messages"], "Step 2 completed"], "step": 2}

    builder = StateGraph(State)
    builder.add_node("step1", step1)
    builder.add_node("step2", step2)
    builder.set_entry_point("step1")
    builder.add_edge("step1", "step2")
    builder.add_edge("step2", END)

    return builder.compile(checkpointer=checkpointer)


def process_a():
    """进程A：创建检查点"""
    print("Process A starting...")

    checkpointer = SqliteSaver("shared_checkpoints.db")
    graph = create_graph(checkpointer)

    initial_state = {"messages": ["Process A started"], "step": 0}
    config = {"configurable": {"thread_id": "shared_thread"}}

    # 执行到中间状态
    result = graph.invoke(initial_state, config=config)

    print(f"Process A completed. Final state: {result}")
    return "shared_thread"


def process_b(thread_id):
    """进程B：恢复检查点"""
    print("Process B starting...")
    time.sleep(2)  # 确保进程A先完成

    if not os.path.exists("shared_checkpoints.db"):
        print("No checkpoint file found")
        return

    checkpointer = SqliteSaver.from_conn("shared_checkpoints.db")
    graph = create_graph(checkpointer)

    config = {"configurable": {"thread_id": thread_id}}

    # 获取当前状态
    snapshot = graph.get_state(config)
    print(f"Process B recovered state: {snapshot.values}")

    # 可以继续执行或读取状态
    if snapshot.next:
        print("Continuing execution from checkpoint...")
        result = graph.invoke(None, config=config)
        print(f"Process B completed: {result}")
    else:
        print("Process already completed")


if __name__ == "__main__":
    # 清理旧数据库
    if os.path.exists("shared_checkpoints.db"):
        os.remove("shared_checkpoints.db")

    # 创建进程
    p1 = multiprocessing.Process(target=process_a)
    p2 = multiprocessing.Process(target=process_b, args=("shared_thread",))

    # 启动进程
    p1.start()
    p2.start()

    # 等待进程完成
    p1.join()
    p2.join()

    print("All processes completed")